Performance: kron-based sparse product kernels, factorization reuse in mldivide, new blkdiag method - #7
Conversation
…use in mldivide, add blkdiag - matvalXmatder / matderXmatval: replace the explicit loop over the inner dimension with a single kron-based sparse product (identities vec(A*dB) = kron(I,A)*vec(dB) and vec(dA*B) = kron(B',I)*vec(dA)). 30-150x faster at kernel level; package-level mtimes for (1000x1000 sparse)*(1000x20) with 100 derivative directions: AD*AD 0.196s -> 0.040s, double*AD 0.125s -> 0.0008s. Results identical. - mldivide: factor the matrix once per call via decomposition() and reuse it for the value solve and all derivative solves. Previously each backslash refactorized: twice for one RHS, once per column for matrix RHS. 13% (1 RHS) to 27% (10 RHS) faster on a 3600x3600 5-point stencil; gains grow with factorization cost. - blkdiag: new method assembling values and stacked derivatives with one sparse() call each. Building block-diagonal AD matrices by repeated concatenation is O(nblocks^2) with a full derivative permutation per step; blkdiag of 30 blocks (40x40, 100 directions) is 68x faster with bit-identical output. - compile_mex_files.m: comment out remaining mex compilation with an explanatory note. On recent MATLAB (tested R2025b) the .m fallbacks are as fast or faster (matdrivXvecval mex measured 4-37x slower than its .m), and a compiled mex silently shadows the .m implementation. All changed paths validated against central finite differences (~1e-10) and on a continuous-time HANK transition-Jacobian application (198x198 Jacobian through backward HJB / forward KFE with sparse solves): Jacobians identical to master to the last bit. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…ests Follow-up to the previous commit, addressing two problems with the kron-based formulation. Memory: kron(I,A) and kron(B',I) materialize ncol*nnz(A) and nrow*nnz(B) entries respectively, regardless of how sparse the derivative payload is. For a dense 2000x2000 value matrix with 100 columns that is 4e8 stored entries (~5 GB), even if the payload has a handful of nonzeros. Both kernels now avoid forming any kron intermediate: - matvalXmatder: the contraction is over the leading index of the column-wise stacking, so reshaping the payload to (ninter x ncol*nderiv) lays every direction side by side and one sparse product handles all of them. Strictly better than kron: same result, O(nnz) peak memory, and faster everywhere measured (dense-A/sparse-payload case 0.0067s -> 0.0001s; HANK-like case 0.0007s -> 0.0003s). - matderXmatval: the contraction is over the trailing index, so the payload is reindexed once (O(nnz)) to move the derivative direction into the row index, one sparse product contracts all directions, and the result is reindexed back. Faster than kron in the wide/dense regimes (5-6x) and within ~30% of it in the moderate-sparse regime, with bounded memory. Complex correctness: neither kernel now takes a transpose at all, so both are correct for complex input. The kron version of matderXmatval used B' (conjugate transpose) where vec(dA*B) = kron(B.',I)*vec(dA) requires the nonconjugate transpose. Note master had an analogous pre-existing bug in both kernels: `Aval = Aval(:)'` conjugates the stored derivative values. Measured max derivative errors for complex input, versus a per-direction reference: master 5.1 and 7.3, kron-with-B' 4.2, both new kernels ~5e-16. Adds test_matrix_kernels.m: checks mtimes (AD*AD, AD*double, double*AD) and mldivide against the per-direction product rule over five shapes, for real and complex inputs, plus zero-payload, single-direction, and blkdiag cases. The complex cases fail on master, which is how the conjugation bug was found. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
|
Follow-up commit (7150b1e) — please review this branch at 7150b1e rather than the initial commit. Two problems with the 1. Memory. Both kernels now avoid forming any
Package-level 2. Complex correctness. Neither kernel takes a transpose now, so both are correct for complex input by construction. My Worth flagging that master has an analogous pre-existing bug in both kernels: Tests. Added Re-validated after the change: finite-difference checks on all touched paths (~1e-10), |
…ling convention The previous note said the mex kernels were merely slower, and described the new matvalXmatder/matderXmatval as kron-based (no longer true after the reformulation). The more important point was missing: matvalXmatder.c expects the derivative of the TRANSPOSED right-hand matrix, as its own header documents and as the old mtimes.m supplied via dertransp. The 2018 .m takes the column-wise stacked derivative instead, and today's mtimes.m calls it that way. Because a compiled mex silently shadows the same-named .m, compiling matvalXmatder.c makes AD matrix-matrix products return wrong derivatives with no warning. Verified: with the column-wise convention the mex is off by O(1); called as mex(A, dertransp(dB, n)) it matches the .m exactly. So the kernel is not buggy in itself -- it is correct through the old API -- but it must not be compiled against the current mtimes. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
|
One correction to something I wrote in the PR description, since it was imprecise in a way that matters. I said the The C header documents it — "To get dB/dx to dB'/dx, you can call dertransp(dB/dx,m) prior to calling matvalXmatder" — and the pre-2018 So the kernel is not buggy in itself; it is correct through the old API. But since a compiled mex silently shadows the same-named Worth flagging for anyone with a stale binary in their working copy: this would have been silently wrong since the 2018 kernel rewrite, not just under this branch. It may be worth deleting the checked-in Also fixed in that commit: the note still described |
1 participant
This PR was prepared at the request of Greg Kaplan, following a performance exploration of the package together with Claude (Anthropic's Claude Code). The motivating application is computing large NxN transition-path Jacobians for continuous-time heterogeneous-agent GE models (backward HJB / forward KFE recursions built from large sparse matrices, sparse linear solves, and embedded nonlinear solvers), where the package is used to differentiate the stacked equilibrium conditions. Profiling that application against current master identified a small number of hot spots; every change below was benchmarked against master and validated to produce identical derivatives.
Changes
1.
private/matvalXmatder.mandprivate/matderXmatval.m: single kron-based sparse productThe 2018 implementations loop over the inner dimension of the matrix product with logical masks over all nonzeros, which is
O(ninter*(nnz(A)+nnz(B))). The identitiesvec(A*dB) = kron(I_ncol, A) * vec(dB)vec(dA*B) = kron(B', I_nrow) * vec(dA)turn each kernel into one sparse matrix product, letting MATLAB's multithreaded sparse BLAS do the work.
Kernel-level timings (R2025b, Apple Silicon;
A5 nnz/row,ncol=20, 100 derivative directions):matvalXmatdermastermatderXmatvalmasterPackage-level effect on
mtimesfor(1000x1000 sparse) * (1000x20)with 100 directions: ADAD 0.196s -> 0.040s (4.9x), doubleAD 0.125s -> 0.0008s (156x), AD*double 0.077s -> 0.040s (1.9x). Derivative fingerprints identical to master.2.
mldivide.m: reuse one factorization per callEach matrix branch previously factorized the same matrix multiple times: twice for a single right-hand side (value solve + derivative solve), and once per column for matrix right-hand sides. The patch creates
dec = decomposition(sparse(...))once per call and reuses it for all solves. Thebinary_exthandling and all branch semantics are unchanged.On a 3600x3600 five-point-stencil matrix with 100 derivative directions: 13% faster with 1 RHS, 27% faster with 10 RHS; identical results. Gains grow with factorization cost relative to triangular-solve cost, i.e. precisely in the large sparse systems the package targets. Requires R2017b+ (
decomposition), which the package effectively already assumes (implicit expansion is used inprivate/valXder.m).3. New method:
blkdiag.mBuilding block-diagonal AD matrices by repeated concatenation, e.g.
is
O(nblocks^2)and permutes the full stacked derivative matrix at every step. In the HANK application above (assembling the intensity matrix from per-income-state blocks inside the HJB loop), this concatenation pattern was 63% of total Jacobian runtime. The newblkdiagmethod assembles the values and the stacked derivatives with onesparse()triplet call each, accepts any mix of myAD and numeric blocks, and matches the concatenation result bit-for-bit. For 30 blocks of 40x40 with 100 directions: 0.225s -> 0.003s (68x).4.
compile_mex_files.m: comment out mex compilation, with an explanatory noteOn recent MATLAB the pure-MATLAB kernels are as fast as or faster than the C mex kernels, which predate multithreaded sparse ops and implicit expansion:
valXdermex is at parity with the one-line.m, and thematdrivXvecvalmex measured 4-37x slower than its 2018.mimplementation (N = 500 to 5000, K = 198). Because a compiled mex silently shadows the.mfile of the same name, running the compile script can make the package slower on modern MATLAB. Compilation lines are kept (commented) for old releases, with a note to delete stale binaries. (Incidentally, thematvalXmatder.cmex produced results differing from both.mimplementations in our tests, so its compilation remaining commented out seems wise.)private/matdrivXvecval.mis deliberately left unchanged: a kron reformulation was faster at moderate sizes but slower at N=5000, so the existing implementation stands.Validation
mtimesADAD / ADdouble / double*AD, all six matrix branches ofmldivide,blkdiagwith mixed AD/numeric blocks) validated against central finite differences: max error ~1e-10.blkdiagoutput compared element-by-element against the concatenation construction: identical values and derivatives.EXAMPLE_AutoDiff_syntax.mruns unchanged.Benchmarks were run on MATLAB R2025b (Apple Silicon, maca64). The kernel changes are pure sparse linear algebra, so relative gains should carry over to other platforms.
🤖 Generated with Claude Code